package org.ws.task.repository;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.item.ExecuteContext;
import org.ws.task.item.TaskInstance;
import org.ws.task.schedule.dispatcher.Dispatcher;
import org.ws.task.schedule.dispatcher.DispatcherFactory;
import org.ws.task.model.TaskConfig;
import org.ws.task.model.TaskItem;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class LeaderContextRepository {

    private static Logger log = LoggerFactory.getLogger(LeaderContextRepository.class.getName());

    private static final LeaderContextRepository instance = new LeaderContextRepository();
    // Key:Map<TaskName,Map<NodeId,Context>>
    private final Map<String, Map<String, ExecuteContext>> taskWorkerNodeContextMap;
    private final ReadWriteLock lock = new ReentrantReadWriteLock();

    private LeaderContextRepository(){
        this.taskWorkerNodeContextMap = Maps.newConcurrentMap();
    }

    public static LeaderContextRepository getInstance(){
        return instance;
    }

    public ExecuteContext getContext(String taskName,String nodeId) {
        int count = 3;
        while (count > 0) {
            try {
                lock.writeLock().lockInterruptibly();
                Map<String, ExecuteContext> nodeContextMap = getNodeContext(taskName);
                if (!nodeContextMap.containsKey(nodeId)) {
                    nodeContextMap.put(nodeId, new ExecuteContext());
                }
                return nodeContextMap.get(nodeId);
            } catch (InterruptedException e) {
                log.error(e.getMessage(), e);
            } finally {
                lock.writeLock().unlock();
                count--;
            }
        }
        return null;
    }

    private Map<String,ExecuteContext> getNodeContext(String taskName){
        if(!this.taskWorkerNodeContextMap.containsKey(taskName)){
            this.taskWorkerNodeContextMap.put(taskName,Maps.newConcurrentMap());
        }
        return this.taskWorkerNodeContextMap.get(taskName);
    }

    public void resize(String taskName) {
        long timeMillis = System.currentTimeMillis();
        try {
            log.info("task[{}] start resize.",taskName);
            lock.writeLock().lockInterruptibly();
            if(!this.taskWorkerNodeContextMap.containsKey(taskName)){
                return;
            }
            // 找到已丢失的nodeId
            Set<String> nodeIdSet = this.taskWorkerNodeContextMap.get(taskName).keySet();
            Set<String> removeIds = Sets.newHashSet();
            for (String id : nodeIdSet) {
                if(!NodeMetadataRepository.getInstance().isMatch(taskName,id)){
                    removeIds.add(id);
                }
            }
            //
            if(!removeIds.isEmpty()){
                List<TaskItem> taskItemList = Lists.newArrayList();
                for (String removeId : removeIds) {
                    ExecuteContext context = getContext(taskName, removeId);
                    while (!context.isEmpty()){
                        taskItemList.add(context.poll());
                    }
                }
                TaskInstance taskInstance = TaskRepository.getInstance().getTaskInstance(taskName);
                TaskConfig.SendModel sendModel = taskInstance.getTaskConfig().getSendModel();
                Dispatcher dispatcher = DispatcherFactory.getInstance().getDispatcher(TaskConfig.DispatchStrategy.POLL,sendModel);
                for (TaskItem taskItem : taskItemList) {
                    dispatcher.dispatch(taskItem,taskName);
                }
            }
        }
        catch (InterruptedException e) {
            log.error(e.getMessage(),e);
        }
        finally {
            lock.writeLock().unlock();
            log.info("task[{}] resize end. {}ms",taskName,(System.currentTimeMillis()-timeMillis));
        }
    }

}
